package com.newbee.model;

import lombok.extern.slf4j.Slf4j;

import javax.annotation.Nonnull;
import java.util.Stack;
import java.util.concurrent.CopyOnWriteArrayList;

@Slf4j
public class TraceNodeStack {

    // 根节点（整条链路的起点）
    private final TraceNode root;

    // 用于同步链路的下一个节点快速关联
    private final Stack<TraceNode> stack;

    public TraceNodeStack(TraceNode root) {
        if (null == root) {
            throw new IllegalArgumentException("TraceNodeStack构造失败,root为空");
        }
        this.root = root;
        this.stack = new Stack<>();
    }

    @Nonnull
    public TraceNode getRoot() {
        return root;
    }

    @Nonnull
    public Stack<TraceNode> getStack() {
        return stack;
    }

    /**
     * 添加节点到栈
     *
     * @param newNode 新节点
     */
    public void pushNodeAndRelatedChildNodes(@Nonnull TraceNode newNode) {

        if (stack.isEmpty()) {
            stack.push(newNode);
            log.info("{} 入栈, 是根节点", newNode.getMethodName());
        } else {
            TraceNode parent = stack.peek();
            stack.push(newNode);
            if (null == parent.getChildren()) {
                parent.setChildren(new CopyOnWriteArrayList<>());
            }
            parent.getChildren().add(newNode);
            log.info("{} 入栈, 关联为 {} 的子节点", newNode.getMethodName(), parent.getMethodName());
        }
    }

    /**
     * 关联为子节点
     *
     * @param asyncTraceNode 异步节点
     */
    public void relatedAsyncChildNodes(@Nonnull TraceNode asyncTraceNode) {
        if (stack.isEmpty()) {
            log.info("异步节点入栈失败");
        } else {
            TraceNode parent = stack.peek();
            if (null == parent.getChildren()) {
                parent.setChildren(new CopyOnWriteArrayList<>());
            }
            parent.getChildren().add(asyncTraceNode);
            log.info("异步节点 入栈, 关联为 {} 的子节点", parent.getMethodName());
        }
    }

    public TraceNode getParent() {
        if (stack.isEmpty()) {
            return null;
        }
        return stack.peek();
    }

    public TraceNode popNode() {
        if (stack.isEmpty()) {
            return null;
        }
        return stack.pop();
    }

    public TraceNode peekNode() {
        if (stack.isEmpty()) {
            return null;
        }
        return stack.peek();
    }
}